from PyQt5.QtCore import QStringListModel
from PyQt5.Qt import QCompleter
from PyQt5.Qt import Qt
from PyQt5.QtWidgets import QWidget
from PyQt5.QtWidgets import QAbstractItemView
import json
from view.FunctionHookView import Ui_FunctionHook
from PyQt5.Qt import QTextDocument, QTextCharFormat, QPalette, QColor, QTextCursor,QFont


class FunctionHook(QWidget, Ui_FunctionHook):
    def __init__(self, frida_manager=None, log_out=None):
        super(FunctionHook, self).__init__()
        self.method_list = list()
        self.method_mapping = dict()
        self.setupUi(self)
        self.setLayout(self.verticalLayout)
        self.frida_manager = frida_manager
        self.log_out = log_out
        self.le_class_name.setText('com.hx.hxsocket.MainActivity')
        self.pb_hook.clicked.connect(self.pb_hook_click)
        self.pb_load_method.clicked.connect(self.pb_load_method_click)
        self.cb_batch.stateChanged.connect(self.batch_click)
        self.cb_overload.stateChanged.connect(self.overload_click)
        self.le_search.textChanged.connect(self.search_changed)

    def search_changed(self):
        text = self.le_search.text()

        document = self.text_edit.document()
        highlight_cursor = QTextCursor(document)
        cursor = QTextCursor(document)
        cursor.beginEditBlock()
        color_init = QTextCharFormat(highlight_cursor.charFormat())
        color_init.setBackground(QColor(0, 0, 0, 0))

        texts = self.get_texts(self.text_edit)
        for item in texts:
            highlight_cursor = document.find(item, highlight_cursor)
            if not highlight_cursor.isNull():
                highlight_cursor.mergeCharFormat(color_init)
        cursor.endEditBlock()

        document = self.text_edit.document()
        highlight_cursor = QTextCursor(document)
        cursor = QTextCursor(document)
        cursor.beginEditBlock()
        color_format = QTextCharFormat(highlight_cursor.charFormat())
        color_format.setBackground(QColor(255, 211, 6))

        while (not highlight_cursor.isNull()) and (not highlight_cursor.atEnd()):
            highlight_cursor = document.find(text, highlight_cursor)
            if not highlight_cursor.isNull():
                highlight_cursor.mergeCharFormat(color_format)
        cursor.endEditBlock()

        list_temp = []
        for item in self.method_list:
            if text in item:
                list_temp.append(item)

        if len(list_temp) > 0:
            list_model = QStringListModel()
            list_model.setStringList(list_temp)
            self.list_view.setModel(list_model)

    def update_class(self, data):
        if data['data'] == 'null':
            return
        completer = QCompleter()
        string_list = QStringListModel()
        string_list.setStringList(data['data'])
        completer.setModel(string_list)
        completer.setCaseSensitivity(Qt.CaseInsensitive)
        completer.setFilterMode(Qt.MatchRecursive)
        completer.setCompletionMode(QCompleter.InlineCompletion)
        self.le_class_name.setCompleter(completer)

    def batch_click(self):
        if self.cb_batch.isChecked():
            self.list_view.setSelectionMode(QAbstractItemView.MultiSelection)
            self.cb_overload.setEnabled(False)
        else:
            self.list_view.setSelectionMode(QAbstractItemView.SingleSelection)
            self.cb_overload.setEnabled(True)

    def overload_click(self):
        if self.cb_overload.isChecked():
            self.cb_batch.setEnabled(False)
        else:
            self.cb_batch.setEnabled(True)

    def update_list_view(self, data):
        try:
            self.text_edit.append('类方法加载完成...')
            if data['data'] == 'null':
                self.method_list.clear()
                list_model = QStringListModel()
                list_model.setStringList(self.method_list)
                self.list_view.setModel(list_model)
                return

            data = json.loads(data['data'])
            self.method_list.clear()
            test = []
            prev = ""
            temp = ""
            for item in data:
                current = self.analysis_method_name(item)
                if current in prev:
                    index = int(prev.split('_')[1]) + 1
                    temp = current + '_' + str(index)
                    test.append(temp)
                else:
                    temp = current + '_0'
                    test.append(temp)

                prev = temp
                self.method_mapping.update({item: temp})
                self.method_list.append(item)

            list_model = QStringListModel()
            list_model.setStringList(self.method_list)

            self.list_view.setModel(list_model)
        except Exception as e:
            self.log_out.insert(self.log_out.exception_msg(self, e))

    def update_text_view(self, data):
        # try:
        self.text_edit.append(str(data))
        # self.text_edit.append('函数 : ' + data['method'])
        # arguments = json.loads(data['arguments'])
        # count = 1
        # for item in arguments:
        #     self.text_edit.append('参数 ' + str(count) + ' : ' + item)
        #     count = count + 1
        #
        # if 'return' in data.keys():
        #     self.text_edit.append('返回值: ' + data['return'])
        #
        # self.text_edit.append('堆栈调用: ')
        # stack = data['stack']
        # self.text_edit.append(stack)
        # except Exception as e:
        #     self.log_out.insert(self.log_out.exception_msg(self, e))

    def insert_class_info(self, data):
        try:
            completer = QCompleter()
            string_list = QStringListModel()
            string_list.setStringList(data['data'])
            completer.setModel(string_list)
            completer.setCaseSensitivity(Qt.CaseInsensitive)
            completer.setFilterMode(Qt.MatchRecursive)
            completer.setCompletionMode(QCompleter.InlineCompletion)
            self.le_class_name.setCompleter(completer)
            self.text_edit.append('类列表加载完成...')
        except Exception as e:
            self.log_out.insert(self.log_out.exception_msg(self, e))

    def pb_load_method_click(self):
        class_name = self.le_class_name.text()
        try:
            if class_name == '':
                return
            self.text_edit.append('正在加载类方法...')
            self.rpc_get_class_methods(class_name)
        except Exception as e:
            self.log_out.insert(self.log_out.exception_msg(self, e))

    def pb_hook_click(self):
        try:
            if not self.frida_manager:
                return

            if len(self.method_list) <= 0:
                return

            methods = dict()
            methods.update({'class_name': self.le_class_name.text()})
            if self.cb_overload.isChecked():
                if len(self.list_view.selectedIndexes()) > 1:
                    raise Exception('多选bug')
                data = self.list_view.currentIndex().data()
                if not data:
                    return

                current = self.analysis_method_name(data)
                count = 0
                for item in self.method_list:
                    if current == self.analysis_method_name(item):
                        count = count + 1
                count = count - 1

                methods.update({'cmd': 'overload'})
                methods.update({'method': current})
                methods.update({'count': count})

            elif self.cb_batch.isChecked():
                methods.update({'cmd': 'batch'})
                temp = []
                for item in self.list_view.selectedIndexes():
                    temp.append(self.method_mapping[item.data()])
                methods.update({'methods': temp})

            else:
                if len(self.list_view.selectedIndexes()) > 1:
                    raise Exception('多选bug')
                data = self.list_view.currentIndex().data()
                if not data:
                    return
                method = self.method_mapping[data]
                methods.update({'cmd': 'general'})
                methods.update({'method': method.split("_")[0]})
                methods.update({'index': method.split("_")[1]})
            self.rpc_get_hook_methods(str(methods))
            self.text_edit.append('函数挂钩成功...')
        except Exception as e:
            self.log_out.insert(self.log_out.exception_msg(self, e))

    def rpc_get_class_methods(self, class_name):
        exports = self.frida_manager.get_exports()
        if exports:
            exports.get_function(class_name)

    def rpc_get_hook_methods(self, methods):
        exports = self.frida_manager.get_exports()
        if exports:
            exports.hook_method(methods)

    def get_overload_methods(self, text):
        methods = []
        for method in self.method_list:
            if text in method:
                methods.append(method)
        return methods

    def on_message(self, message, data):
        if message['type'] == 'send':
            text = message['payload']
            js = json.loads(text)
            if js['name'] == 'functions':
                self.update_list_view(js)
            elif js['name'] == 'hook':
                self.update_text_view(js)
            elif js['name'] == 'loadclass':
                self.update_class(js)
        else:
            print(message)

    @staticmethod
    def analysis_method_name(text):
        text = text.split('(')
        text = text[0]
        start = text.rfind('.')
        return text[start + 1:]

    @staticmethod
    def get_texts(edit):
        texts = edit.toPlainText()
        texts = texts.split('\n')
        texts = [i for i in texts if i != '']
        return texts

